import streamlit as st
import numpy as np
import torch
from torch import nn

# 加载模型
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(LSTMModel, self).__init__()
        self.lstm1 = nn.LSTM(input_size, hidden_size, num_layers=2, batch_first=True, dropout=0.2)
        self.lstm2 = nn.LSTM(hidden_size, hidden_size, num_layers=1, batch_first=True, dropout=0.2)
        self.fc1 = nn.Linear(hidden_size, 64)
        self.fc2 = nn.Linear(64, output_size)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x, _ = self.lstm1(x)  # 第一层LSTM
        x, _ = self.lstm2(x)  # 第二层LSTM
        x = self.fc1(x[:, -1, :])  # 取最后一个时间步的输出并通过全连接层
        x = self.fc2(x)  # 输出层
        return self.sigmoid(x)  # 使用Sigmoid激活函数

# 初始化模型
model = LSTMModel(input_size=8, hidden_size=64, output_size=1)
model.load_state_dict(torch.load('diabetes_lstm_model.pth'))
model.eval()

# 设置页面标题和布局
st.title('糖尿病预测')

# 特征输入
features = ['怀孕次数', '葡萄糖浓度', '血压', '皮肤厚度', '胰岛素', 'BMI', '糖尿病谱系函数', '年龄']
inputs = []
for feature in features:
    if feature == '怀孕次数':
        value = st.number_input(feature, value=0, step=1, format="%d")
    else:
        value = st.number_input(feature, value=0.0, step=0.1)
    inputs.append(value)

# 预测按钮
if st.button('预测'):
    try:
        inputs = np.array(inputs).reshape(1, 1, 8)
        inputs = torch.tensor(inputs, dtype=torch.float32)

        output = model(inputs)
        predicted = (output.squeeze() > 0.5).item()
        if predicted:
            result_text = '糖尿病'
            result_color = 'red'
        else:
            result_text = '非糖尿病'
            result_color = 'green'

        st.markdown(f'<p style="color:{result_color}; font-size:24px; font-weight:bold;">预测结果: {result_text}</p>', unsafe_allow_html=True)
    except Exception as e:
        st.error(f'发生错误: {str(e)}')
